//----------------------------------------------------------------------------
// M A R I T I M E  R E S E A R C H  I N S T I T U T E  N E T H E R L A N D S
//----------------------------------------------------------------------------
// Copyright (C) 2012 - MARIN - All rights reserved - http://www.marin.nl
//----------------------------------------------------------------------------
// Program    : mermaid
// Module     : vbm
// File       : prec_kernels.h
// Author     : M.A. de Jong
//----------------------------------------------------------------------------

#ifndef PREC_KERNELS_H
#define	PREC_KERNELS_H

namespace rrbsolver
{

template <class T>
__global__ void kernel_prec1(Grid g)
{
    int ld = g.nx;

    int u_b1 = BORDER_WIDTH2 + g.cx + bx * Bx + tx;
    int v_b1 = BORDER_WIDTH + by * By + ty;
    int u_b2 = BORDER_WIDTH + bx * Bx + tx;
    int v_b2 = BORDER_WIDTH2 + g.cy + by * By + ty;

    int loc_b1 = ld * v_b1 + u_b1;
    int loc_b2 = ld * v_b2 + u_b2;

    g.cc[loc_b1] = g.cc[loc_b1] + 
                   g.ne[loc_b1] +
                   g.nw[loc_b1] +
                   g.se[loc_b1] + 
                   g.sw[loc_b1];

    g.cc[loc_b2] = g.cc[loc_b2] +
                   g.ne[loc_b2] +
                   g.nw[loc_b2] +
                   g.se[loc_b2] + 
                   g.sw[loc_b2];
}

template <class T>
__global__ void kernel_prec2(Grid g)
{
    int ld = g.nx;

    int u_r1 = BORDER_WIDTH + bx * Bx + tx;
    int v_r1 = BORDER_WIDTH + by * By + ty;
    int u_r2 = BORDER_WIDTH2 + g.cx + bx * Bx + tx;
    int v_r2 = BORDER_WIDTH2 + g.cy + by * By + ty;
    int u_b1 = BORDER_WIDTH2 + g.cx + bx * Bx + tx;
    int v_b1 = BORDER_WIDTH + by * By + ty;
    int u_b2 = BORDER_WIDTH + bx * Bx + tx;
    int v_b2 = BORDER_WIDTH2 + g.cy + by * By + ty;

    int loc_r1 = ld * v_r1 + u_r1;
    int loc_r2 = ld * v_r2 + u_r2;
    int loc_b1 = ld * v_b1 + u_b1;
    int loc_b2 = ld * v_b2 + u_b2;

    T sum_cc = 0;
    T sum_ne = 0;
    T sum_nw = 0;
    T sum_nn = 0;
    T sum_se = 0;
    T sum_sw = 0;
    T sum_ss = 0;
    T sum_ee = 0;
    T sum_ww = 0;

    const T val_nn_b1 = __FETCH_NN(u_b1, v_b1, ld);
    const T val_ee_b1 = __FETCH_EE(u_b1, v_b1, ld);
    const T val_ss_b1 = __FETCH_SS(u_b1, v_b1, ld);
    const T val_ww_b1 = __FETCH_WW(u_b1, v_b1, ld);

    const T val_nn_b2 = __FETCH_NN(u_b2, v_b2, ld);
    const T val_ee_b2 = __FETCH_EE(u_b2, v_b2, ld);
    const T val_ss_b2 = __FETCH_SS(u_b2, v_b2, ld);
    const T val_ww_b2 = __FETCH_WW(u_b2, v_b2, ld);

    const T val_cc_b1 = __FETCH_CC(u_b1, v_b1, ld);
    const T val_cc_b2 = __FETCH_CC(u_b2, v_b2, ld);

    T a, b;

    // r1 nodes:
    a = 1 / val_cc_b2;
    b = g.nn[loc_r1] * a;
    sum_ne  = g.ne[loc_r1] - b * val_ee_b2;
    sum_cc  = g.cc[loc_r1] - b * val_ss_b2;
    sum_nw  = g.nw[loc_r1] - b * val_ww_b2;
    sum_nn  =              - b * val_nn_b2;

    a = 1 / __FETCH_CC(u_b2, v_b2-1, ld);
    b = g.ss[loc_r1] * a;
    sum_se  = g.se[loc_r1] - b * g.ee[loc_b2 - ld];
    sum_sw  = g.sw[loc_r1] - b * g.ww[loc_b2 - ld];
    sum_cc -=                b * g.nn[loc_b2 - ld];
    sum_ss  =              - b * g.ss[loc_b2 - ld];
 
    a = 1 / val_cc_b1;
    b = g.ee[loc_r1] * a;
    sum_ne -=                b * val_nn_b1;
    sum_se -=                b * val_ss_b1;
    sum_cc -=                b * val_ww_b1;
    sum_ee  =              - b * val_ee_b1;

    a = 1 / __FETCH_CC(u_b1-1, v_b1, ld);
    b = g.ww[loc_r1] * a;
    sum_nw -=                b * __FETCH_NN(u_b1-1, v_b1, ld);
    sum_sw -=                b * __FETCH_SS(u_b1-1, v_b1, ld);
    sum_cc -=                b * __FETCH_EE(u_b1-1, v_b1, ld);
    sum_ww  =              - b * __FETCH_WW(u_b1-1, v_b1, ld);

    g.cc[loc_r1] = sum_cc + sum_nn + sum_ee + sum_ss + sum_ww; 
    g.nn[loc_r1] = sum_ne;
    g.ee[loc_r1] = sum_se;
    g.ww[loc_r1] = sum_nw;
    g.ss[loc_r1] = sum_sw;
    g.ne[loc_r1] = sum_ne;
    g.se[loc_r1] = sum_se;
    g.sw[loc_r1] = sum_sw;
    g.nw[loc_r1] = sum_nw;

    // r2 nodes:
    a = 1 / __FETCH_CC(u_b1, v_b1+1, ld);
    b = g.nn[loc_r2] * a;
    sum_ne  = g.ne[loc_r2] - b * g.ee[loc_b1 + ld];
    sum_cc  = g.cc[loc_r2] - b * g.ss[loc_b1 + ld];
    sum_nw  = g.nw[loc_r2] - b * g.ww[loc_b1 + ld];
    sum_nn  =              - b * g.nn[loc_b1 + ld];

    a = 1 / val_cc_b1;
    b = g.ss[loc_r2] * a;
    sum_se  = g.se[loc_r2] - b * val_ee_b1;
    sum_sw  = g.sw[loc_r2] - b * val_ww_b1;
    sum_cc -=                b * val_nn_b1;
    sum_ss  =              - b * val_ss_b1;

    a = 1 / __FETCH_CC(u_b2+1, v_b2, ld);
    b = g.ee[loc_r2] * a;   
    sum_ne -=                b * __FETCH_NN(u_b2+1, v_b2, ld);
    sum_se -=                b * __FETCH_SS(u_b2+1, v_b2, ld);
    sum_cc -=                b * __FETCH_WW(u_b2+1, v_b2, ld);
    sum_ee  =              - b * __FETCH_EE(u_b2+1, v_b2, ld);

    a = 1 / val_cc_b2;
    b = g.ww[loc_r2] * a;
    sum_nw -=                b * val_nn_b2;
    sum_sw -=                b * val_ss_b2;
    sum_cc -=                b * val_ee_b2;
    sum_ww  =              - b * val_ww_b2;

    g.cc[loc_r2] = sum_cc;
    g.nn[loc_r2] = sum_nn;
    g.ss[loc_r2] = sum_ss;
    g.ee[loc_r2] = sum_ee;
    g.ww[loc_r2] = sum_ww;
    g.ne[loc_r2] = sum_ne;
    g.se[loc_r2] = sum_se;
    g.sw[loc_r2] = sum_sw;
    g.nw[loc_r2] = sum_nw;
}

template <class T>
__global__ void kernel_prec3(Grid g)
{
    int ld = g.nx;

    int u_r1 = BORDER_WIDTH + bx * Bx + tx;
    int v_r1 = BORDER_WIDTH + by * By + ty;
    int u_r2 = BORDER_WIDTH2 + g.cx + bx * Bx + tx;
    int v_r2 = BORDER_WIDTH2 + g.cy + by * By + ty;

    int loc_r2 = ld * v_r2 + u_r2;

    T sum_cc = 0;
    T sum_ne = 0;
    T sum_nw = 0;
    T sum_nn = 0;
    T sum_se = 0;
    T sum_sw = 0;
    T sum_ss = 0;
    T sum_ee = 0;
    T sum_ww = 0;

    T a, b;

    const T val_cc_r1 = 1 / __FETCH_CC(u_r1, v_r1, ld);

    a = 1 / __FETCH_CC(u_r1+1, v_r1+1, ld);
    b = g.ne[loc_r2] * a;
    sum_nn  = g.nn[loc_r2] - b * __FETCH_NW(u_r1+1, v_r1+1, ld);
    sum_ee  = g.ee[loc_r2] - b * __FETCH_SE(u_r1+1, v_r1+1, ld);
    sum_cc  = g.cc[loc_r2] - b * __FETCH_SW(u_r1+1, v_r1+1, ld);
    sum_ne  =              - b * __FETCH_NE(u_r1+1, v_r1+1, ld);

    a = 1 / __FETCH_CC(u_r1, v_r1+1, ld);
    b = g.nw[loc_r2] * a;
    sum_nn -=                b * __FETCH_NE(u_r1, v_r1+1, ld);
    sum_ww  = g.ww[loc_r2] - b * __FETCH_SW(u_r1, v_r1+1, ld);
    sum_cc -=                b * __FETCH_SE(u_r1, v_r1+1, ld);
    sum_nw  =              - b * __FETCH_NW(u_r1, v_r1+1, ld);

    a = val_cc_r1;
    b = g.sw[loc_r2] * a;
    sum_ss  = g.ss[loc_r2] - b * __FETCH_SE(u_r1, v_r1, ld);
    sum_ww -=                b * __FETCH_NW(u_r1, v_r1, ld);
    sum_cc -=                b * __FETCH_NE(u_r1, v_r1, ld);
    sum_sw  =              - b * __FETCH_SW(u_r1, v_r1, ld);

    a = 1 / __FETCH_CC(u_r1+1, v_r1, ld);
    b = g.se[loc_r2] * a;
    sum_ss -=                b * __FETCH_SW(u_r1+1, v_r1, ld);
    sum_ee -=                b * __FETCH_NE(u_r1+1, v_r1, ld);
    sum_cc -=                b * __FETCH_NW(u_r1+1, v_r1, ld);
    sum_se  =              - b * __FETCH_SE(u_r1+1, v_r1, ld);

    g.cc[loc_r2] = sum_cc;
    g.ne[loc_r2] = sum_ne;
    g.nw[loc_r2] = sum_nw;
    g.nn[loc_r2] = sum_nn;
    g.se[loc_r2] = sum_se;
    g.sw[loc_r2] = sum_sw;
    g.ss[loc_r2] = sum_ss;
    g.ee[loc_r2] = sum_ee;
    g.ww[loc_r2] = sum_ww;
}

template <class T>
__global__ void kernel_prec4(Grid g)
{
    int ld = g.nx;

    int u_r1 = BORDER_WIDTH + bx * Bx + tx;
    int v_r1 = BORDER_WIDTH + by * By + ty;
    int u_r2 = BORDER_WIDTH2 + g.cx + bx * Bx + tx;
    int v_r2 = BORDER_WIDTH2 + g.cy + by * By + ty;
    int u_b1 = BORDER_WIDTH2 + g.cx + bx * Bx + tx;
    int v_b1 = BORDER_WIDTH + by * By + ty;
    int u_b2 = BORDER_WIDTH + bx * Bx + tx;
    int v_b2 = BORDER_WIDTH2 + g.cy + by * By + ty;

    int loc_r1 = ld * v_r1 + u_r1;
    int loc_r2 = ld * v_r2 + u_r2;
    int loc_b1 = ld * v_b1 + u_b1;
    int loc_b2 = ld * v_b2 + u_b2;

    T fac;

    // r1 nodes:
    fac = 1 / g.cc[loc_r1];
    g.nn[loc_r1] *= fac;
    g.ss[loc_r1] *= fac;
    g.ee[loc_r1] *= fac;
    g.ww[loc_r1] *= fac;
    g.cc[loc_r1]  = fac;

    // r2 nodes:
    fac = 1 / g.cc[loc_r2];
    g.nn[loc_r2] *= fac;
    g.ss[loc_r2] *= fac;
    g.ee[loc_r2] *= fac;
    g.ww[loc_r2] *= fac;
    g.cc[loc_r2]  = fac;

    // b1 nodes:
    fac = 1 / g.cc[loc_b1];
    g.nn[loc_b1] *= fac;
    g.ss[loc_b1] *= fac;
    g.ee[loc_b1] *= fac;
    g.ww[loc_b1] *= fac;
    g.cc[loc_b1]  = fac;

    // b2 nodes:
    fac = 1 / g.cc[loc_b2];
    g.nn[loc_b2] *= fac;
    g.ss[loc_b2] *= fac;
    g.ee[loc_b2] *= fac;
    g.ww[loc_b2] *= fac;
    g.cc[loc_b2]  = fac;
}

template <class T>
__global__ void kernel_prec5(Grid g)
{
    int ld = g.nx;

    int v_b1 = BORDER_WIDTH + by * By + ty;
    int u_b1 = BORDER_WIDTH2 + g.cx + bx * Bx + tx;
    int v_b2 = BORDER_WIDTH2 + g.cy + by * By + ty;
    int u_b2 = BORDER_WIDTH + bx * Bx + tx;

    int loc_b1 = ld * v_b1 + u_b1;
    int loc_b2 = ld * v_b2 + u_b2;

    const T c_b1 = g.cc[loc_b1];
    const T c_b2 = g.cc[loc_b2];

    g.cc[loc_b1] = 1 / c_b1;
    g.cc[loc_b2] = 1 / c_b2;
}

template <class T>
__global__ void kernel_precfinal(Grid g)
{
    int ld = g.nx;

    int loc;

    int st = 1;
    int tw = 2;
    int u, v;

    T a, b;

    for (int nt = DIM_COMPUTE_BLOCK / 2; nt >= 1; nt >>= 1) 
    {
        // phase 1: for all black points
        if (tx < nt) {
            u = (st - 1) + tw * tx;
            v = (tw - 1) + tw * ty;
        } else {
            u = (tw - 1) + tw * (tx - nt);
            v = (st - 1) + tw * ty;
        }
        __syncthreads();
        loc = ld * (BORDER_WIDTH2 + DIM_COMPUTE_BLOCK + v) 
                  + BORDER_WIDTH2 + DIM_COMPUTE_BLOCK + u;
        if (ty < nt && tx < nt << 1) {
            g.cc[loc] = g.cc[loc] + 
                        g.ne[loc] + 
                        g.nw[loc] +
                        g.se[loc] + 
                        g.sw[loc]; 
        }
        __syncthreads();

        // phase 2: for all red points
        if (tx < nt) {
            u = (st - 1) + tw * tx;
            v = (st - 1) + tw * ty;
        } else {
            u = (tw - 1) + tw * (tx - nt);
            v = (tw - 1) + tw * ty;
        }
        loc = ld * (BORDER_WIDTH2 + DIM_COMPUTE_BLOCK + v) 
                  + BORDER_WIDTH2 + DIM_COMPUTE_BLOCK + u;

        if (ty < nt && tx < nt << 1) {
            a = 1 / g.cc[loc + st * ld];
            b = g.nn[loc] * a;
            g.ne[loc] -=   b * g.ee[loc + st * ld];
            g.cc[loc] -=   b * g.ss[loc + st * ld];
            g.nw[loc] -=   b * g.ww[loc + st * ld];
            g.nn[loc]  = - b * g.nn[loc + st * ld];

            a = 1 / g.cc[loc - st * ld];
            b = g.ss[loc] * a;
            g.se[loc] -=   b * g.ee[loc - st * ld];
            g.sw[loc] -=   b * g.ww[loc - st * ld];
            g.cc[loc] -=   b * g.nn[loc - st * ld];
            g.ss[loc]  = - b * g.ss[loc - st * ld];

            a = 1 / g.cc[loc + st];
            b = g.ee[loc] * a;
            g.ne[loc] -=   b * g.nn[loc + st];
            g.se[loc] -=   b * g.ss[loc + st];
            g.cc[loc] -=   b * g.ww[loc + st];
            g.ee[loc]  = - b * g.ee[loc + st];

            a = 1 / g.cc[loc - st];
            b = g.ww[loc] * a;
            g.nw[loc] -=   b * g.nn[loc - st];
            g.sw[loc] -=   b * g.ss[loc - st];
            g.cc[loc] -=   b * g.ee[loc - st];
            g.ww[loc]  = - b * g.ww[loc - st];
        }
        __syncthreads();

        // phase 3: for red points not in next level
        if (ty < nt && tx < nt) {
            g.cc[loc] = g.cc[loc] + 
                        g.nn[loc] + 
                        g.ee[loc] + 
                        g.ss[loc] + 
                        g.ww[loc]; 
            g.nn[loc] = g.ne[loc];
            g.ee[loc] = g.se[loc];
            g.ww[loc] = g.nw[loc];
            g.ss[loc] = g.sw[loc];          
        }
        __syncthreads();

        // phase 4: for red points in next level
        u = (tw - 1) + tw * tx;
        v = (tw - 1) + tw * ty;
        loc = ld * (BORDER_WIDTH2 + DIM_COMPUTE_BLOCK + v) 
                  + BORDER_WIDTH2 + DIM_COMPUTE_BLOCK + u;
        if (ty < nt && tx < nt) {
            a = 1 / g.cc[loc + st * ld + st];
            b = g.ne[loc] * a;
            g.nn[loc] -=   b * g.nw[loc + st * ld + st];
            g.ee[loc] -=   b * g.se[loc + st * ld + st];
            g.cc[loc] -=   b * g.sw[loc + st * ld + st];
            g.ne[loc]  = - b * g.ne[loc + st * ld + st];

            a = 1 / g.cc[loc + st * ld - st];
            b = g.nw[loc] * a;
            g.nn[loc] -=   b * g.ne[loc + st * ld - st];
            g.ww[loc] -=   b * g.sw[loc + st * ld - st];
            g.cc[loc] -=   b * g.se[loc + st * ld - st];
            g.nw[loc]  = - b * g.nw[loc + st * ld - st];

            a = 1 / g.cc[loc - st * ld - st];
            b = g.sw[loc] * a;
            g.ss[loc] -=   b * g.se[loc - st * ld - st];
            g.ww[loc] -=   b * g.nw[loc - st * ld - st];
            g.cc[loc] -=   b * g.ne[loc - st * ld - st];
            g.sw[loc]  = - b * g.sw[loc - st * ld - st];

            a = 1 / g.cc[loc - st * ld + st];
            b = g.se[loc] * a;
            g.ss[loc] -=   b * g.sw[loc - st * ld + st];
            g.ee[loc] -=   b * g.ne[loc - st * ld + st];
            g.cc[loc] -=   b * g.nw[loc - st * ld + st];
            g.se[loc]  = - b * g.se[loc - st * ld + st];
        }
        __syncthreads();

        st = tw;
        tw *= 2;
    }
}

} // namespace rrbsolver

#endif // PREC_KERNELS_H
